// This file is part of Spectrum Editor Called SECS.
//
// Spectrum Editor Called SECS is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// Spectrum Editor Called SECS is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with Spectrum Editor Called SECS.  If not, see <http://www.gnu.org/licenses/>.

struct Sinusoid {
  double time, frequency, amplitude, phase, probability;
  bool is_already_tracked;

  Sinusoid(double _time, double _frequency, double _amplitude, double _phase) {
    time = _time;
    frequency = _frequency;
    amplitude = _amplitude;
    phase = _phase;
    probability = 1.0;
    is_already_tracked = false;
  } 

  Signal get_spectrum(int sample_rate, int window_size, int fft_size, std::vector<double> window_function) {
    Signal normalized_signal(sample_rate);
    for (int i = 0; i < window_size; i++)
      normalized_signal.data.push_back(sin(frequency * 2 * M_PI * i / sample_rate));
    Signal normalized_spectrum = normalized_signal.window(window_function).fft(window_size, fft_size).normalize().to_polar();
    for (int i = 0; i < normalized_spectrum.data.size(); i++)
      normalized_spectrum.data.at(i) *= amplitude;
    return normalized_spectrum;
  }
};

struct Window {
  int index, sample_rate, window_size, fft_size;
  std::vector<Sinusoid> partials;

  Window(int _index, int _window_size, int _fft_size, std::vector<double> window_function, Signal signal) {
    index = _index;
    sample_rate = signal.sample_rate;
    window_size = _window_size;
    fft_size = _fft_size;
    Signal spectrum = signal.window(window_function).fft(window_size, fft_size).to_polar();

    double root_mean_square = 0.3 * spectrum.root_mean_square(); // TODO amplitude threshold for recognizing partial?
    int maximum_bin = spectrum.find_maximum();
    while (real(spectrum[maximum_bin]) > root_mean_square) {
      partials.push_back(Sinusoid(index * window_size / 2.0, bin_to_frequency(maximum_bin), real(spectrum[maximum_bin]), imag(spectrum[maximum_bin])));
      spectrum.subtract_signal(partials[partials.size() - 1].get_spectrum(signal.sample_rate, window_size, fft_size, window_function));
      maximum_bin = spectrum.find_maximum();
    }
    for (int i = 0; i < partials.size(); i++)
      partials[i].amplitude = 10 * log10(partials[i].amplitude);
    std::sort(partials.begin(), partials.end(), compare_frequencies);
  }

  double bin_to_frequency(int bin) {
    return (double)bin * sample_rate / fft_size; 
  }

  static bool compare_frequencies(Sinusoid left, Sinusoid right) {
    return left.frequency < right.frequency;
  }

  static bool compare_amplitudes(Sinusoid left, Sinusoid right) {
    return left.amplitude < right.amplitude;
  }

  void print() {
    for (int i = 0; i < partials.size(); i++)
      std::cout << partials[i].frequency << "  " << partials[i].amplitude << std::endl;
    std::cout << std::endl;
  }

};

struct Partial {
  std::vector<Sinusoid*> points;
  bool is_unterminated;

  Partial(Sinusoid * first_point) {
    points.push_back(first_point);
    is_unterminated = true;
  }
};

struct Analysis {
  Signal signal;
  int window_size, fft_size;
  std::vector<double> window;
  std::vector<Window> windows;

  Analysis(int _window_size, Signal _signal) {
    window_size = _window_size;
    fft_size = 4 * (int)pow(2, (int)(log2(2 * window_size) + 1));
    signal = _signal;
    window = hann_window(window_size);
  }

  std::vector<double> hann_window(int window_size) {
    std::vector<double> window;
    for (int i = 0; i < window_size; i++)
      window.push_back(pow(sin(M_PI * i / (window_size - 1)), 2));
    return window;
  }

  std::vector<double> blackman_harris_window(int window_size) {
    std::vector<double> window;
    double a0 = 0.35875;
    double a1 = 0.48829;
    double a2 = 0.14128;
    double a3 = 0.01168;
    double W = 2 * M_PI / (window_size - 1);
    for (int i = 0; i < window_size; i++)
      window.push_back(a0 - a1 * cos(W * i) + a2 * cos(2 * W * i) - a3 * cos(3 * W * i));
    return window;
  }

  double threshold_curve(double frequency) {
    double b = 0.0075;
    double a_l = -24.0;
    double a_r = 32.0;
    return a_l - a_r + a_r * pow(b, frequency / 20.0);
  }

  void analyze(int width, int height) {
    double maximum_frequency = 600.0;
    double death_threshold = -96.0;

    std::vector<double> window_function = hann_window(window_size);

    std::vector< std::vector<bool> > plot;
    plot.resize(width);
    for (int i = 0; i < width; i++)
      plot.at(i).resize(height, 0);

    int number_of_windows = (signal.data.size() / (window_size / 2)); 

    for (int i = 0; i < number_of_windows - 1; i++)
      windows.push_back(Window(i, window_size, fft_size, window_function, signal.get_segment(i * window_size / 2, i *  window_size / 2 + window_size)));

    std::vector<Partial> partials;

    for (int i = 0; i < number_of_windows - 1; i++) {
      std::vector<double> errors(partials.size());
      for (int j = 0; j < partials.size(); j++) {
        if (i * window_size / 2 - partials[j].points[partials[j].points.size() - 1]->time > 3.0 * window_size)
          partials[j].is_unterminated = true;
        else {
          double predicted_frequency, predicted_amplitude;
          if (partials[j].points.size() == 1) {
            predicted_frequency = partials[j].points[0]->frequency;
            predicted_amplitude = partials[j].points[0]->amplitude;
          }
          else {
            Sinusoid* previous_peak = partials[j].points[partials[j].points.size() - 1];
            Sinusoid* preprevious_peak = partials[j].points[partials[j].points.size() - 2];
            double frequency_slope = (previous_peak->frequency - preprevious_peak->frequency) / (previous_peak->time - preprevious_peak->time);
            double amplitude_slope = (previous_peak->amplitude - preprevious_peak->amplitude) / (previous_peak->time - preprevious_peak->time);
            predicted_frequency = frequency_slope * (i * window_size / 2 - previous_peak->time) + previous_peak->frequency;
            predicted_amplitude = amplitude_slope * (i * window_size / 2 - previous_peak->time) + previous_peak->amplitude;
          }
          int index_of_best_partial;
          int error_of_best_partial;
          for (int k = 0; k < windows[i].partials.size(); k++) {
            double alpha = 0.8;
            double frequency_error = 12.0 * log10(windows[i].partials[k].frequency / predicted_frequency);
            double amplitude_error = alpha * 20.0 * log10(windows[i].partials[k].amplitude / predicted_amplitude);
            double error = sqrt(frequency_error * frequency_error + amplitude_error * amplitude_error);
            if (k == 0) {
              index_of_best_partial = 0;
              error_of_best_partial = error;
            }
            else if (error < error_of_best_partial) {
              index_of_best_partial = k;
              error_of_best_partial = error;
            }
          }
          if (abs(predicted_frequency - windows[i].partials[index_of_best_partial].frequency) < 1.5 * signal.sample_rate / window_size) {
            if (!windows[i].partials[index_of_best_partial].is_already_tracked) {
              partials[j].points.push_back(&windows[i].partials[index_of_best_partial]);
              windows[i].partials[index_of_best_partial].is_already_tracked = true;
              errors[j] = error_of_best_partial;
            }
            else {
              int old_track;
              for (int k = 0; k < partials.size(); k++) {
                if (partials[k].points[partials[k].points.size() - 1] == &windows[i].partials[index_of_best_partial]) {
                  old_track = k;
                  break;
                }
              }
              double new_error = abs(partials[old_track].points[partials[old_track].points.size() - 1]->frequency - predicted_frequency);
              if (new_error < errors[old_track]) {
                partials[old_track].points.erase(partials[old_track].points.begin() + partials[old_track].points.size() - 1);
                partials[j].points.push_back(&windows[i].partials[index_of_best_partial]);
                errors[j] = error_of_best_partial;
              }
            }
          }
        }
      }
     
      for (int j = 0; j < windows[i].partials.size(); j++) {
        if (!windows[i].partials[j].is_already_tracked)
          partials.push_back(Partial(&windows[i].partials[j]));
      }
    }    

    
    for (int i = 0; i < partials.size(); i++) {
      if (partials[i].points.size() > 3) {
        for (int j = 0; j < partials[i].points.size(); j++)
          std::cout << (partials[i].points[j]->time / signal.sample_rate) << " " <<  partials[i].points[j]->frequency << std::endl;
        std::cout << std::endl;
      }
    }
  }
};
